import numpy as np
import random
from math import sqrt, pi
from scipy.spatial import KDTree

class Node:
    def __init__(self, config):
        self.config = np.array(config)  # 6维关节角度向量
        self.parent = None
        self.cost = 0.0
        self.children = []

    def distance_to(self, other):
        """计算6维空间欧氏距离（兼容Node对象和列表输入）"""
        if isinstance(other, Node):
            target = other.config
        else:  # 处理列表输入
            target = np.array(other)
        return sqrt(np.sum((self.config - target) ** 2))

def sample_random_config(goal_config, bias=0.1):
    """带目标偏置的随机采样（6维空间）"""
    if random.random() < bias:
        return goal_config
    return [random.uniform(-pi, pi) for _ in range(6)]

def steer(from_node, to_config, step_size=pi/18):
    """关节空间路径扩展"""
    direction = np.array(to_config) - from_node.config
    norm = np.linalg.norm(direction)
    
    if norm < 1e-6:  # 零向量处理
        return from_node
    
    step = min(step_size, norm) * direction / norm
    new_config = from_node.config + step
    return Node(new_config)

def rrt_star_6d(start_config, goal_config, max_iter=2000, 
                step_size=pi/18, radius=1.0):
    # 初始化树结构
    start_node = Node(start_config)
    goal_node = Node(goal_config)
    nodes = [start_node]
    goal_threshold = step_size * 2  # 目标区域阈值
    
    for _ in range(max_iter):
        # 阶段1：随机采样
        rand_config = sample_random_config(goal_config)
        
        # 阶段2：KDTree加速最近邻查询（网页3方法）
        kdtree = KDTree([node.config for node in nodes])
        _, indices = kdtree.query(rand_config)
        nearest = nodes[indices[0]] if isinstance(indices, np.ndarray) else nodes[indices]
        
        # 阶段3：路径扩展
        new_node = steer(nearest, rand_config, step_size)
        
        # 阶段4：父节点优化（网页5的NearToNew逻辑）
        near_indices = kdtree.query_ball_point(new_node.config, radius)
        near_nodes = [nodes[i] for i in near_indices]
        
        # 选择最优父节点
        min_cost = float('inf')
        best_parent = nearest
        for node in near_nodes:
            cost = node.cost + node.distance_to(new_node)
            if cost < min_cost:  # 已移除碰撞检测
                min_cost = cost
                best_parent = node
        
        # 更新节点关系
        new_node.parent = best_parent
        new_node.cost = min_cost
        best_parent.children.append(new_node)
        
        # 阶段5：重布线优化（网页2的rewire逻辑）
        for node in near_nodes:
            if node != best_parent:
                new_cost = new_node.cost + new_node.distance_to(node)
                if new_cost < node.cost:  # 已移除碰撞检测
                    node.parent.children.remove(node)
                    node.parent = new_node
                    node.cost = new_cost
                    new_node.children.append(node)
        
        nodes.append(new_node)
        
        # 目标检查（网页1的终止条件）
        if new_node.distance_to(goal_config) < goal_threshold:
            final_node = Node(goal_config)
            final_node.parent = new_node
            final_node.cost = new_node.cost + new_node.distance_to(goal_node)
            nodes.append(final_node)
            print(f"✅ 目标在{len(nodes)}个节点时找到")
            break
    
    # 路径回溯（网页4的路径重构方法）
    path = []
    current = nodes[-1] if nodes[-1].distance_to(goal_config) < goal_threshold else None
    while current:
        path.append(current.config.tolist())
        current = current.parent
    path.reverse()

    return path

# ========================
# 示例用法
# ========================
if __name__ == '__main__':
    # 6维关节角度配置（单位：弧度）
    start = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
    goal = [pi/2, 0, pi/2, 0, pi/2, 0]
    
    # 执行路径规划
    path = rrt_star_6d(start, goal)
    
    # 输出结果
    print("\n规划路径节点序列：")
    for i, config in enumerate(path):
        print(f"Step {i}: {np.round(config, 3)}")
    
    print(f"\n规划路径总节点数：",path)